#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: bao
# @Date:   2017-01-16 11:59:29
# @Last Modified by:   bao
# @Last Modified time: 2017-05-14 11:05:36

from datetime import datetime

from CoreSim.Event import *
from CoreSim.Scheduler import *
from CoreSim.Unit import *


class SchedulerFair(Scheduler):
    """ Flow scheduler with fairness bandwidth allocation"""

    def __init__(self, host_num, switches_num):
        Scheduler.__init__(self)

        self.host_num = host_num
        self.switches_num = switches_num
        self.hosts_num = host_num * switches_num

        self.flowsDict = {}  # flowId -> flow
        self.runningFlowsDict = {}  # running flowId -> flow

    def UpdateFlow(self, curFlow, eventType, eventQueue, curTime):
        """
        Fairness bandwidth allocation
        Find related flows and update their flowBw, transTime, endTime.
        """
        if eventType == FlowArrival:
            curFlow.updateTime = curTime
            self.runningFlows.append(curFlow)
            self.flowsDict[curFlow.flowId] = curFlow
            self.runningFlowsDict[curFlow.flowId] = curFlow

            pathInLink = curFlow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                if curFlow.flowId not in link.flowIds:
                    link.flowIds.append(curFlow.flowId)

        elif eventType == FlowFinished:

            self.UpdateEnergy([curFlow], curTime)

            # update finished flow status
            curFlow.finishTime = curTime
            pathInLink = curFlow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                link.flowIds.remove(curFlow.flowId)

            self.runningFlows.remove(curFlow)
            self.finishedFlows.append(curFlow)
            self.runningFlowsDict.pop(curFlow.flowId)
            self.flowsDict.pop(curFlow.flowId)

            if len(self.finishedFlows) % 100 == 0:
                print len(self.finishedFlows), curTime

        elif eventType == UpdateTopology:
            for flow in self.runningFlows:
                pathInLink = flow.pathLinkIds
                for linkId in pathInLink:
                    link = self.Links[linkId]
                    if flow.flowId not in link.flowIds:
                        link.flowIds.append(flow.flowId)

        current_flows = list(self.runningFlows)
        current_links = list(self.Links)

        # remove flows without routing
        for flow in self.runningFlows:
            if not flow.pathLinkIds:
                current_flows.remove(flow)

        residual_cap = {}
        residual_flowIds = {}
        for link_id in self.Links:
            link = self.Links[link_id]
            residual_cap[link_id] = link.linkCap
            residual_flowIds[link_id] = list(link.flowIds)
            self.Links[link_id].residualCap = link.linkCap

        self.UpdateEnergy(self.runningFlows, curTime)

        for flow in self.runningFlows:
            # update flow status
            flow.remainSize -= flow.bw * (curTime - flow.updateTime)
            flow.remainTime -= curTime - flow.updateTime
            flow.updateTime = curTime
            flow.bw = 0.0

            # update flow and eventQueue
            if flow.finishEvent != 0:
                flow.finishEvent.cancelled = True

        while current_flows:  # multi iterations
            if not current_links:
                break

            bottleneck_rate = link.linkCap
            linksToRemove = []
            for link_id in current_links:
                if len(residual_flowIds[link_id]) != 0:
                    if bottleneck_rate >= residual_cap[link_id] / len(residual_flowIds[link_id]):
                        bottleneck_rate = residual_cap[link_id] / len(residual_flowIds[link_id])
                else:
                    linksToRemove.append(link_id)

            # delete the empty links
            for link_id in linksToRemove:
                current_links.remove(link_id)

            bottleneck_rate = max([min([bottleneck_rate, link.linkCap]), 0.001])
            # update flow bandwidth
            for flow in current_flows:
                if flow.pathLinkIds:
                    flow.bw += bottleneck_rate

            # delete bottleneck links and update related flows
            linksToRemove = []
            for link_id in current_links:
                residual_cap[link_id] -= bottleneck_rate * len(residual_flowIds[link_id])
                self.Links[link_id].residualCap = residual_cap[link_id]

                if residual_cap[link_id] <= 1.0e-20:
                    linksToRemove.append(link_id)
                    flowIds = list(residual_flowIds[link_id])

                    for flow_id in flowIds:
                        flow = self.runningFlowsDict[flow_id]
                        pathInLink = flow.pathLinkIds
                        for flow_link_id in pathInLink:
                            residual_flowIds[flow_link_id].remove(flow_id)
                        current_flows.remove(flow)
                        # update the flow status

                        flow.finishTime = curTime + flow.remainSize / flow.bw + PROPAGATION_DELAY * len(
                            flow.pathLinkIds)
                        event = FlowFinishedEvent(flow.finishTime, flow)
                        flow.finishEvent = event
                        eventQueue.push(event, event.time)

            # delete the bottleneck links
            for link_id in linksToRemove:
                current_links.remove(link_id)

        del current_flows
        del current_links
        del residual_cap
        del residual_flowIds

    def UpdateEnergy(self, flows, curTime):
        """
        # calculate energy
        """
        for flow in flows:
            pathInLink = flow.pathLinkIds
            energy = 0
            for (srcId, dstId) in pathInLink:
                srcNodeRole = GetNodeRole(srcId, self.hosts_num, self.switches_num, 0)
                dstNodeRole = GetNodeRole(dstId, self.hosts_num, self.switches_num, 0)

                if srcNodeRole == Host and dstNodeRole == ToR:
                    energy += PortPower * 2 + EPSPower + OCSPower
                elif srcNodeRole == ToR and dstNodeRole == ToR:
                    energy += PortPower * 2 + EPSPower + OCSPower
                elif srcNodeRole == ToR and dstNodeRole == Host:
                    energy += PortPower * 2
            flow.energy += energy * flow.bw * (curTime - flow.updateTime)

    def CalBandwidth(self):
        """
        Fairness bandwidth allocation
        Find related flows and update their flowBw, transTime, endTime.
        """
        startTime = datetime.now()
        for curFlow in self.flows:
            curFlow.bw = 0.0
            self.flowsDict[curFlow.flowId] = curFlow
            pathInLink = curFlow.pathLinkIds
            for linkId in pathInLink:
                link = self.Links[linkId]
                if curFlow.flowId not in link.flowIds:
                    link.flowIds.append(curFlow.flowId)

        current_flows = list(self.flows)
        current_links = list(self.Links)
        residual_cap = {}
        residual_flowIds = {}
        for link_id in self.Links:
            link = self.Links[link_id]
            residual_cap[link_id] = link.linkCap
            residual_flowIds[link_id] = list(link.flowIds)

        while current_flows:  # multi iterations
            # find bottleneck rate
            bottleneck_rate = link.linkCap
            linksToRemove = []
            for link_id in current_links:
                if len(residual_flowIds[link_id]) != 0:
                    if bottleneck_rate >= residual_cap[link_id] / len(residual_flowIds[link_id]):
                        bottleneck_rate = residual_cap[link_id] / len(residual_flowIds[link_id])
                else:
                    linksToRemove.append(link_id)
            # delete the empty links
            for link_id in linksToRemove:
                current_links.remove(link_id)

            # update flow bandwidth
            bottleneck_rate = max([min([bottleneck_rate, link.linkCap]), 0.001])
            for flow in current_flows:
                flow.bw += bottleneck_rate

            # delete bottleneck links and update related flows
            linksToRemove = []
            for link_id in current_links:
                residual_cap[link_id] -= bottleneck_rate * len(residual_flowIds[link_id])
                if residual_cap[link_id] <= 1.0e-20:
                    linksToRemove.append(link_id)
                    flowIds = list(residual_flowIds[link_id])
                    for flow_id in flowIds:
                        flow = self.flowsDict[flow_id]
                        pathInLink = flow.pathLinkIds

                        for flow_link_id in pathInLink:
                            residual_flowIds[flow_link_id].remove(flow_id)
                        current_flows.remove(flow)

            # delete the bottleneck links
            for link_id in linksToRemove:
                current_links.remove(link_id)

        del current_flows
        del current_links
        del residual_cap
        del residual_flowIds
        endTime = datetime.now()
        timeDiff = endTime - startTime
        print ("CalBandwidth ended. Execution time: %s \n" % timeDiff)
